package com.zhanghp.class035;

import java.util.PriorityQueue;

/**
 * 快速获得数据流的中位数的结构<br/>
 * 测试链接 : https://leetcode.cn/problems/find-median-from-data-stream/
 * <p>
 * 大根堆 & 小跟堆实现
 * </p>
 * <p>
 *     <ul>实现原理
 *     <li>
 *          1. 2个堆的目的：如果偶数，则直接拿2个堆的首，奇数：直接拿size大的堆的首节点
 *          <br/>
 *          2. 2个堆的size相差<=1
 *          <br/>
 *          3. 若放入时，相差==2，则把size大的首节点弹出放到size小的，平衡一下
 *          <br/>
 *          4. 放入规则：大根堆里，只放小于等于首节点的值，否则放小跟堆
 *     </li>
 *     <li>大根堆 和 小跟堆相等，则2个堆的首节点peek，相加除2</li>
 *     <li>大根堆 和 小跟堆，差1，则返回size大的根堆的首节点的值</li>
 *     </ul>
 * </p>
 *
 * @author zhanghp
 * @since 2024/4/16 11:28
 */
public class Code05_MedianFinder {

    static class MedianFinder {

        public PriorityQueue<Integer> maxHeap;

        public PriorityQueue<Integer> minHeap;


        public MedianFinder() {
            maxHeap = new PriorityQueue<>((a, b) -> b - a);
            minHeap = new PriorityQueue<>((a, b) -> a - b);
        }

        public void addNum(int num) {
            if (maxHeap.isEmpty() || maxHeap.peek() >= num) {
                maxHeap.add(num);
            } else {
                minHeap.add(num);
            }
            // 平衡大小
            balance();
        }

        public double findMedian() {
            // 规则： 在调用 findMedian 之前，数据结构中至少有一个元素
            // 所以不用判断size为0的情况
            if (maxHeap.size() == minHeap.size()) {
                return (double)(maxHeap.peek() + minHeap.peek()) / 2;
            }else {
                return maxHeap.size() > minHeap.size() ? maxHeap.peek() : minHeap.peek();
            }
        }

        private void balance() {
            if (Math.abs(maxHeap.size() - minHeap.size()) == 2) {
                if (maxHeap.size() > minHeap.size()) {
                    minHeap.offer(maxHeap.poll());
                } else {
                    maxHeap.offer(minHeap.poll());
                }
            }
        }
    }
}
